from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.prompts import ChatPromptTemplate
from langchain.agents import AgentExecutor, create_tool_calling_agent
import datetime
from langchain.chat_models import init_chat_model

# 设置环境变量 (替换为你的实际API密钥)
import os

# os.environ["OPENAI_API_KEY"] = os.getenv("DEEPSEEK_API_KEY")
# base_url = os.getenv("DEEPSEEK_BASE_URL")
os.environ["OPENAI_API_KEY"] = os.getenv("BAOCLOUD_DS_API_KEY")
base_url = os.getenv("BAOCLOUD_DS_BASE_URL")


# 定义自定义工具1：计算器
@tool
def calculator(expression: str) -> float:
    """执行数学计算。支持加减乘除和括号。示例: "(2 + 3) * 4" """
    print("call calculator tool")
    try:
        # 安全评估数学表达式
        return eval(expression, {"__builtins__": None}, {})
    except Exception as e:
        return f"计算错误: {str(e)}"


# 定义自定义工具2：获取当前时间
@tool
def get_current_time(timezone: str = "UTC") -> str:
    """获取指定时区的当前时间。时区示例: 'Asia/Shanghai'"""
    print("call get current time tool")
    try:
        # 创建时区感知的当前时间
        tz = datetime.timezone.utc if timezone.upper() == "UTC" else datetime.datetime.now().astimezone().tzinfo
        return datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S %Z")
    except Exception as e:
        return f"时间获取失败: {str(e)}"


@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息。如果用户询问天气，必须使用此工具。参数city是城市名称，例如：'北京'、'上海'。"""
    print("call get weather tool")
    try:
        # 模拟获取天气信息
        return f"{city}的天气是晴天，温度为25度。"
    except Exception as e:
        return f"天气获取失败: {str(e)}"


# 初始化聊天模型 (使用支持tool calling的模型)
# model = ChatOpenAI(model="deepseek-chat", temperature=0, base_url=base_url)
model = init_chat_model(model="DeepSeek-V3", model_provider="openai", base_url=base_url)

# 创建工具列表
tools = [calculator, get_current_time, get_weather]

# 创建提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system",
     "你是一个有用的助手。请使用工具精确回答问题。当用户提问天气怎么样时，你必须使用get_weather工具，且用户必须指明城市，如未指明，则返回并提示用户说明他想问的是哪个城市"),
    ("placeholder", "{chat_history}"),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}")
])

# 创建工具调用代理
agent = create_tool_calling_agent(model, tools, prompt)

# 创建代理执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,  # 显示详细执行过程
    handle_parsing_errors=True  # 自动处理解析错误
)

# 运行代理示例
if __name__ == "__main__":
    res = agent_executor.invoke({"input": "今天上海天气怎么样？"})
    print(f"答案：{res['output']}")
